import os
os.environ["OMP_NUM_THREADS"]="1"

import pandas as pd
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from sklearn.metrics import calinski_harabasz_score     # CH系数
from sklearn.metrics import silhouette_score            # SC指标


# 一、使用评估方法求出最优的超参数
def kmeans_evaluate(type):
    # 1- 准备数据
    df = pd.read_csv("customers.csv",encoding="UTF-8")
    x = df.iloc[:,3:]

    # 2- 肘方法：循环训练
    value_list = []

    start=2
    if type=="SSE":
        start=1

    for i in range(start, 101):
        # 2.1- 模型训练
        model = KMeans(n_clusters=i)
        y_predict = model.fit_predict(x)

        # 2.2- 获得评价指标
        if type == "SSE":
            # SSE指标
            value_list.append(model.inertia_)
        elif type == "CH":
            # CH轮廓系数指标
            value_list.append(calinski_harabasz_score(x, y_predict))
        else:
            # SC指标
            value_list.append(silhouette_score(x, y_predict))

    # 3- 绘制肘方法的图形：横轴是n_clusters超参数；纵轴是评价指标
    plt.figure(figsize=(20, 20), dpi=150)
    plt.plot(range(start, 101), value_list, c="r")
    plt.ylabel(type)
    plt.xticks(range(0, 101, 2))
    plt.grid()
    plt.show()


# 二、模型构建、学习
def ml():
    # 1.准备数据
    df = pd.read_csv("customers.csv", encoding="UTF-8")

    # 1.1 获取到特征值
    x_data = df.iloc[:, 3:]

    # 2.模型构建
    model = KMeans(n_clusters=5)

    # 3.模型预测
    y_pre = model.fit_predict(x_data)

    # 4.模型评估
    print("SSE指标", model.inertia_)
    print("CH", calinski_harabasz_score(x_data, y_pre))
    print("SC", silhouette_score(x_data, y_pre))

    # 5.绘制图形
    # 5.1- 获得质心的信息
    centers = model.cluster_centers_
    print(centers)

    # 5.2- 样本数据绘制成散点
    plt.scatter(x_data.values[y_pre == 0, 0], x_data.values[y_pre == 0, 1], c="red", s=100)
    plt.scatter(x_data.values[y_pre == 1, 0], x_data.values[y_pre == 1, 1], c="blue", s=100)
    plt.scatter(x_data.values[y_pre == 2, 0], x_data.values[y_pre == 2, 1], c="green", s=100)
    plt.scatter(x_data.values[y_pre == 3, 0], x_data.values[y_pre == 3, 1], c="yellow", s=100)
    plt.scatter(x_data.values[y_pre == 4, 0], x_data.values[y_pre == 4, 1], c="pink", s=100)

    # 5.3- 绘制质心
    plt.scatter(centers[:, 0], centers[:, 1], c="black", s=300)
    plt.show()


if __name__ == '__main__':
    # 根据这三个指标来确定出K的最优参数
    kmeans_evaluate("SSE")
    kmeans_evaluate("CH")
    kmeans_evaluate("SC")
    ml()